import pickle as pkl
import networkx as nx
from matplotlib import pyplot as plt


datasetname = 'rpt_rt_tree_graph'  # 'crg_gnp_random_graph', 'rpt_rt_tree_graph','rc_bg_graph'
graphname = 'rpt_rt'  # 'crg_gnp_p' p=0.2~0.9. 'rpt_rt','rc_bg'
# datasetname = 'rpt_rt_tree_graph' # 'crg_gnp_random_graph', 'rpt_rt_tree_graph','rc_bg_graph'
# graphname = 'rpt_rt'  # 'crg_gnp_p' p=0.2~0.9. 'rpt_rt','rc_bg'

# data location
# original_data_path = 'E:/2023\chap5_data_results/' \
#                      'data/{}/'.format(datasetname)
# data information
original_data_path = '../marldata20231011/{}/'\
        .format(datasetname)

nx_original_data_file = '{}.pkl'.format(graphname)
nx_graph_data_location = original_data_path + nx_original_data_file

# transformed_data_filename = '{}_struc2vec_transform.pkl'.format(graphname)
transformed_data_filename = '{}_transform20231011.pkl'.format(graphname)
transformed_pyg_graph_data_location = original_data_path + transformed_data_filename


# to select
flag = 0

if flag==0:
    my_data = pkl.load(open(nx_graph_data_location, 'rb'))

    id = 32         #135:9 746:2,496 #32 #746, 779:3, 849:3, 521:4, 351:4, 770:9, 713:5, 793:2, 571:2
    print(my_data[id])
    # adj = nx.from_scipy_sparse_array(my_data[0]['adj'])
    # print(adj.nodes)
    # # g = nx.from_scipy_sparse_matrix(adj])
    # # data = from_networkx(g)
    # print(adj)
    #spectral_layout,nx.spring_layout,random_layout,shell_layout,circular_layout
    nodelist = my_data[id]['g'].nodes
    print(my_data[id]['dist'])
    print('nodelist=',nodelist)
    pos = nx.spring_layout(my_data[id]['g'])
    nx.draw(my_data[id]['g'],
            node_size=100,
            node_color='y',
            node_shape='o',
            with_labels=True,
            font_size=8,
            pos=pos,
            font_color='r')
    plt.show()
   # print(my_data[0])
elif flag==1:
    my_data = pkl.load(open(transformed_pyg_graph_data_location, 'rb'))
    id = 1
    print(my_data[id])
    print(my_data[id].x)
    print(my_data[id].y[0])




